import os
from service.llm_service import LLM_Service
from prompts.prompts import LAWS_EXTRACTION_SYS_PROMPT, LAWS_EXTRACTION_PROMPT, LAWS_MERGE_PROMPT
from utils.utils import read_file, write_file, write_dict_to_json
import pandas as pd
import json
import sys
from pathlib import Path

ROOT_DIR = Path(__file__).resolve().parent
print(ROOT_DIR)

sys.path.append(str(ROOT_DIR))  # 添加项目根目录到模块搜索路径

"""
该文件用于提取关于法律文件的问询规则
"""

llm_service = LLM_Service()

class Laws_Service:
    def export_query_about_laws(self, query_df, problem_df, project_df, save_dir, save_name):
        """
        导出和法律相关的问询
        :return:
        """
        data = []
        for index in range(query_df.shape[0]):
            row = query_df.iloc[index, :]
            query_content = row['query_content']
            if "《" in query_content and "》" in query_content:
                query_id = row['query_id']
                problem_id = row['problem_id']
                project_id = problem_df[problem_df['problem_id'] == problem_id]['project_id']
                round = problem_df[problem_df['problem_id'] == problem_id]['round'].to_list()[0]
                company_name = project_df[project_df['project_id'] == project_id]['company_name'].to_list()[0]
                problem_title = problem_df[problem_df['problem_id'] == problem_id]['problem_title'].to_list()[0]
                query_order = row['query_order']
                temp = [company_name, round, problem_title, query_order, query_content]
                data.append(temp)
        df = pd.DataFrame(data, columns=['company_name', 'round', 'problem_title', 'query_order', 'query_content'])
        save_path = os.path.join(save_dir, save_name)
        df.to_csv(save_path, index=False)

    def batch_generator_query_rules(self, category_content,df_laws, save_dir, batch_size=50):
        """
        批量生成法规的问询
        :return:
        """
        temp = []
        for index in range(df_laws.shape[0]):
            row = df_laws.iloc[index, :]
            if (index + 1) % batch_size == 0:
                question_content = LAWS_EXTRACTION_PROMPT.format(
                    question='\n'.join(temp),
                    ref_category=category_content)
                response = llm_service.call_with_messages(LAWS_EXTRACTION_SYS_PROMPT, question_content)
                print(response)
                save_path = os.path.join(save_dir, str(index + 1 - batch_size) + "-" + str(index + 1) + '_laws.json')
                write_file(save_path, response)
                temp = []
            else:
                temp.append(row['query_content'])

    def merge_laws_query_rules(self, folder_path):
        """
        合并所有关于法律的规则
        :param map_path:
        :return:
        """
        res = dict()
        for filename in os.listdir(folder_path):
            if filename.endswith('.json'):
                file_path = os.path.join(folder_path, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        if isinstance(data, dict):
                            for key in data:
                                if key not in res:
                                    res[key] = data[key]
                                else:
                                    res[key].extend(data[key])
                            print(f"成功加载 {filename}")
                except (json.JSONDecodeError, FileNotFoundError) as e:
                    print(f'无法正确加载{filename}:{e}')
                    continue
        return res

    def count_laws_query_rules(self, file_path):
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                for key in data:
                    print("key: %s, len(rules): %s" % (key, str(len(data[key]))))
        except (json.JSONDecodeError, FileNotFoundError) as e:
            print(f'无法正确加载{file_path}:{e}')
            return


    def drop_duplicate_laws_query_rules(self, file_path, save_path, batch_size=80):
        """
        半自动去重法律文件的问询规则
        """
        # 读取json文件
        try:
            if not os.path.exists(save_path):
                os.makedirs(save_path, exist_ok=True)
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                for key, value in data.items():
                    # print("key: %s, len(rules): %s" % (key, str(len(value))))
                    question_content = LAWS_MERGE_PROMPT.format(category=key, query_rules_about_laws=str(value))
                    print(question_content)
                    # 保存query
                    write_file(os.path.join(save_path, key+'_query_about_laws_prompt.txt'), question_content)
                    if len(value) <= batch_size: # 直接合并
                        response = llm_service.call_with_messages(LAWS_EXTRACTION_SYS_PROMPT, question_content)
                        print('%s 类型规则，合并后的问询规则如下：\n' % key, response)
                        write_file(os.path.join(save_path, key + '_query_about_laws.json'), response)
                    else: # 如果规则数大于100，每100进行合并
                        for start in range(0, len(value), batch_size):
                            end = min(start + batch_size, len(value))
                            question_content = LAWS_MERGE_PROMPT.format(category=key, query_rules_about_laws=str(value[start:end]))
                            print(question_content)
                            response = llm_service.call_with_messages(LAWS_EXTRACTION_SYS_PROMPT, question_content)
                            write_file(os.path.join(save_path, key + '_query_about_laws_' + str(start) + '-' + str(end) + '.json'), response)
                merged_res = self.merge_laws_query_rules(save_path)
                write_dict_to_json(os.path.join(save_path, 'laws_query_rules_merged.json'), merged_res)
        except (json.JSONDecodeError, FileNotFoundError) as e:
            print(f'无法正确加载{file_path}:{e}')
            return


if __name__ == '__main__':
    laws_service = Laws_Service()
    laws_service.count_laws_query_rules(r"D:\workspace\ipo2\produce_data\laws_query_rules.json")

    input_path = r"D:\workspace\ipo2\produce_data\laws_query_rules.json"
    output_dir = r"D:\workspace\ipo2\produce_data\laws_query"
    laws_service.drop_duplicate_laws_query_rules(input_path, output_dir)

